# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags
import logging
import pandas
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import coin.strategy.mm.tool.archive_base as abase
from coin.strategy.mm.tool.signal_plot_base import dump_book, BookCompress


def launch(trading_date, baseflags, fsflags, subsymbols):
  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.size'] = 9
  plt.rcParams['legend.fontsize'] = 9
  plt.rcParams['xtick.labelsize'] = 9
  plt.rcParams['ytick.labelsize'] = 9

  products = abase.subsymbols_to_products(subsymbols, baseflags)
  start_time, end_time = abase.get_time_range(baseflags)

  feedsub_request = abase.products_to_feedsub_request_with_api_override(
      products, fsflags.api_override, fsflags.feed_machine)

  product_to_sub_request = {}
  for sub_request, products in feedsub_request.gen_iter():
    for product in products:
      product_to_sub_request[product] = sub_request

  for subsymbol, product in zip(subsymbols, products):
    bbos = dump_book(product=product,
                     compress=BookCompress.ALL,
                     trading_date=trading_date,
                     time_range=baseflags.time_range,
                     feed_sub_request=feedsub_request)
    df = pandas.DataFrame(bbos)
    interval_ns = flags.FLAGS.busy_interval_length_min * 60 * 1e9
    nfeed = int(interval_ns / df.timestamp.diff().median())
    idx = df.timestamp.diff(nfeed).idxmin()
    pick_df = df.iloc[(idx - nfeed):idx, :]
    tindex = pandas.DatetimeIndex(pick_df.timestamp)
    plt.plot(tindex,
             pick_df['bid0_price'].fillna(method='ffill'),
             'k-',
             tindex,
             pick_df['ask0_price'].fillna(method='ffill'),
             'k-',
             tindex,
             pick_df['buy_trade_price'],
             'gx',
             tindex,
             pick_df['sell_trade_price'],
             'rx',
             lw=0.2,
             markersize=0.2,
             drawstyle='steps-post')
    plt.xticks(rotation=90)
    plt.minorticks_on()
    plt.grid(lw=0.5, which='major')
    plt.grid(lw=0.2, which='minor')
    filename = f"{flags.FLAGS.out_dir}/feed_busy_{trading_date.strftime('%Y%m%d')}_{subsymbol}.png"
    plt.title(filename)
    print(filename)
    plt.savefig(filename)
    plt.close()


def main(argv):
  jobs = []
  trading_dates = abase.get_trading_dates(flags.FLAGS.trading_date)
  for trading_date in trading_dates:
    flags.FLAGS.trading_date = trading_date.strftime("%Y%m%d")
    baseflags = abase.get_base_flags()
    fsflags = abase.get_feed_subscription_flags()
    jobtuple = (launch, (trading_date, baseflags, fsflags, flags.FLAGS.subsymbols.split(",")))
    if len(trading_dates) == 1:
      jobtuple[0](*jobtuple[1])
    else:
      jobs.append(jobtuple)

  jobs_iter = iter(jobs)
  futures = {}
  from concurrent.futures import as_completed, ProcessPoolExecutor
  with ProcessPoolExecutor(max_workers=flags.FLAGS.num_cpu) as executor:
    for job in jobs_iter:
      future = executor.submit(job[0], *job[1])
      futures[future] = 1
      if len(futures) >= flags.FLAGS.num_cpu:
        for future in as_completed(futures):
          del futures[future]
    for future in futures:
      future.result()


# pick busiest 1 minute interval, usage:
# ./pyrunner python/coin/strategy/mm/tool/feed_busy.py \
# --trading_date=20200401-20200421 \
# --feed_machine=feed-05.ap-northeast-1.aws \
# --subsymbols=Binance:Futures:BTC-USDT.PERPETUAL \
# --busy_interval_length_min=1 --time_range=0-24 --num_cpu=10
if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  flags.FLAGS.feed_machine = 'feed-05.ap-northeast-1.aws'
  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_string('out_dir', 'pic', 'Output directory.')
  flags.DEFINE_float('busy_interval_length_min', 1, 'busy_interval_length_min')
  flags.DEFINE_integer('num_cpu', 12, '')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
